import numpy as np
import random
import math
import copy

''' 种群初始化函数 '''

def initial(pop, dim, ub, lb):
    X = np.zeros([pop, dim])
    for i in range(pop):
        for j in range(dim):
            X[i, j] = random.random()*(ub[j] - lb[j]) + lb[j]

    return X, lb, ub


'''边界检查函数'''


def BorderCheck(X, ub, lb, pop, dim):
    for i in range(pop):
        for j in range(dim):
            if X[i, j] > ub[j]:
                X[i, j] = ub[j]
            elif X[i, j] < lb[j]:
                X[i, j] = lb[j]
    return X


'''计算适应度函数'''


def CaculateFitness(X, fun):
    pop = X.shape[0]
    fitness = np.zeros([pop, 1])
    for i in range(pop):
        fitness[i] = fun(X[i, :])
    return fitness


'''适应度排序'''


def SortFitness(Fit):
    fitness = np.sort(Fit, axis=0)
    index = np.argsort(Fit, axis=0)
    return fitness, index


'''根据适应度对位置进行排序'''


def SortPosition(X, index):
    Xnew = np.zeros(X.shape)
    for i in range(X.shape[0]):
        Xnew[i, :] = X[index[i], :]
    return Xnew


'''鱼鹰优化算法'''


def OOA(pop, dim, lb, ub, MaxIter, fun):
    X, lb, ub = initial(pop, dim, ub, lb)  # 初始化种群
    fitness = CaculateFitness(X, fun)  # 计算适应度值
    fitness, sortIndex = SortFitness(fitness)  # 对适应度值排序
    X = SortPosition(X, sortIndex)  # 种群排序
    GbestScore = copy.copy(fitness[0])
    GbestPositon = np.zeros([1, dim])
    GbestPositon[0, :] = copy.copy(X[0, :])
    Curve = np.zeros([MaxIter, 1])
    for t in range(MaxIter):
        #print("第"+str(t)+"次迭代")
        # 阶段一:定位和捕鱼
        for i in range(pop):
            selectedFish = np.zeros([1,dim])
            if i<2:
                selectedFish[0,:]=copy.copy(GbestPositon[0,:])
            else:
                if np.random.random()<0.5:
                    selectedFish[0,:]=copy.copy(GbestPositon[0,:])
                else:
                    k=np.random.randint(int(i-1))
                    selectedFish[0,:]=copy.copy(X[k,:])
            I=int(round(1+np.random.random())) #{1,2}中的随机数
            X_P1=np.zeros([1,dim])
            X_P1[0,:]=X[i,:]+np.random.random()*(selectedFish-I*X[i,:])
            for j in range(dim):
                if  X_P1[0,j]>ub[j]:
                    X_P1[0,j]=ub[j]
                if X_P1[0,j]<lb[j]:
                    X_P1[0,j]=lb[j]
            fitTemp = fun(X_P1[0,:])
            if fitTemp<fitness[i]:
                fitness[i]=fitTemp
                X[i,:]=copy.copy(X_P1[0,:])
            # 阶段二：把鱼带到合适位置

            X_P1[0,:]=X[i,:]+(lb.T+np.random.random()*(ub.T-lb.T))/(t+1)
            for j in range(dim):
                if  X_P1[0,j]>ub[j]:
                    X_P1[0,j]=ub[j]
                if X_P1[0,j]<lb[j]:
                    X_P1[0,j]=lb[j]
            fitTemp = fun(X_P1[0,:])
            if fitTemp<fitness[i]:
                fitness[i]=fitTemp
                X[i,:]=copy.copy(X_P1[0,:])

        X = BorderCheck(X, ub, lb, pop, dim)  # 边界检测
        fitness = CaculateFitness(X, fun)  # 计算适应度值
        fitness, sortIndex = SortFitness(fitness)  # 对适应度值排序
        X = SortPosition(X, sortIndex)  # 种群排序
        if(fitness[0] <= GbestScore):  # 更新全局最优
            GbestScore = copy.copy(fitness[0])
            GbestPositon[0, :] = copy.copy(X[0, :])
        Curve[t] = GbestScore
        print(f"第 {str(t+1)} 次迭代：最佳参数 =  {int(GbestPositon[0,0]),round((GbestPositon[0, 1]), 5),round((GbestPositon[0, 2]), 5),int(GbestPositon[0,3]),int(GbestPositon[0,4])},最佳适应度 = {GbestScore}")
    return GbestScore, GbestPositon, Curve
